package hbbft

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"math/big"
)

// 简化的椭圆曲线点结构
type Point struct {
	X *big.Int
	Y *big.Int
}

// 阈值加密公钥
type TPKEPublicKey struct {
	N int           // 总节点数
	K int           // 阈值
	VK *Point       // 验证密钥
	VKs []*Point    // 各节点的验证密钥
}

// 阈值加密私钥
type TPKEPrivateKey struct {
	*TPKEPublicKey
	I int           // 节点索引
	SK *big.Int     // 私钥份额
}

// 加密密文
type Ciphertext struct {
	U *Point  // G1中的点
	V []byte  // 32字节
	W *Point  // G2中的点
}

// 解密份额
type DecryptionShare struct {
	U_i *Point
}

// 简化的椭圆曲线操作
func (p *Point) Add(q *Point) *Point {
	// 简化的椭圆曲线点加法
	// 实际实现需要使用椭圆曲线库
	return &Point{
		X: new(big.Int).Add(p.X, q.X),
		Y: new(big.Int).Add(p.Y, q.Y),
	}
}

func (p *Point) ScalarMult(k *big.Int) *Point {
	// 简化的标量乘法
	// 实际实现需要使用椭圆曲线库
	return &Point{
		X: new(big.Int).Mul(p.X, k),
		Y: new(big.Int).Mul(p.Y, k),
	}
}

// 生成阈值加密密钥对
func GenerateTPKEKeys(n, k int) (*TPKEPublicKey, []*TPKEPrivateKey, error) {
	if k > n {
		return nil, nil, errors.New("threshold k cannot be greater than n")
	}

	// 生成随机多项式系数
	secret := new(big.Int)
	secret.SetBytes(randBytes(32))
	
	coefficients := make([]*big.Int, k)
	coefficients[0] = secret
	
	for i := 1; i < k; i++ {
		coefficients[i] = new(big.Int)
		coefficients[i].SetBytes(randBytes(32))
	}

	// 计算各节点的私钥份额
	SKs := make([]*big.Int, n)
	for i := 0; i < n; i++ {
		SKs[i] = evaluatePolynomial(coefficients, big.NewInt(int64(i+1)))
	}

	// 生成验证密钥
	VK := &Point{
		X: new(big.Int).Set(secret), // 简化的G2点
		Y: new(big.Int).Set(secret),
	}
	
	VKs := make([]*Point, n)
	for i := 0; i < n; i++ {
		VKs[i] = &Point{
			X: new(big.Int).Set(SKs[i]),
			Y: new(big.Int).Set(SKs[i]),
		}
	}

	pk := &TPKEPublicKey{
		N:   n,
		K:   k,
		VK:  VK,
		VKs: VKs,
	}

	// 生成私钥
	privKeys := make([]*TPKEPrivateKey, n)
	for i := 0; i < n; i++ {
		privKeys[i] = &TPKEPrivateKey{
			TPKEPublicKey: pk,
			I:            i,
			SK:           SKs[i],
		}
	}

	return pk, privKeys, nil
}

// 多项式求值
func evaluatePolynomial(coefficients []*big.Int, x *big.Int) *big.Int {
	result := new(big.Int)
	xx := big.NewInt(1)
	
	for i, coeff := range coefficients {
		term := new(big.Int).Mul(coeff, xx)
		result.Add(result, term)
		xx.Mul(xx, x)
		_ = i // 避免未使用变量警告
	}
	
	return result
}

// 拉格朗日插值系数
func (pk *TPKEPublicKey) lagrange(S []int, j int) *big.Int {
	if len(S) != pk.K {
		panic("invalid subset size")
	}
	
	num := big.NewInt(1)
	den := big.NewInt(1)
	
	for _, jj := range S {
		if jj != j {
			// num *= (0 - jj - 1)
			term := new(big.Int).Sub(big.NewInt(0), big.NewInt(int64(jj+1)))
			num.Mul(num, term)
			
			// den *= (j - jj)
			term = new(big.Int).Sub(big.NewInt(int64(j)), big.NewInt(int64(jj)))
			den.Mul(den, term)
		}
	}
	
	// 简化的除法，实际实现需要模逆元
	return new(big.Int).Div(num, den)
}

// 加密
func (pk *TPKEPublicKey) Encrypt(m []byte) (*Ciphertext, error) {
	if len(m) != 32 {
		return nil, errors.New("message must be 32 bytes")
	}
	
	// 生成随机数r
	r := new(big.Int)
	r.SetBytes(randBytes(32))
	
	// U = g1^r (简化)
	U := &Point{
		X: new(big.Int).Set(r),
		Y: new(big.Int).Set(r),
	}
	
	// V = m XOR hash(VK^r)
	VK_r := pk.VK.ScalarMult(r)
	hash := sha256.Sum256(VK_r.X.Bytes())
	V := make([]byte, 32)
	for i := 0; i < 32; i++ {
		V[i] = m[i] ^ hash[i]
	}
	
	// W = hash(U, V)^r (简化)
	hashInput := append(U.X.Bytes(), V...)
	hashW := sha256.Sum256(hashInput)
	W := &Point{
		X: new(big.Int).SetBytes(hashW[:16]),
		Y: new(big.Int).SetBytes(hashW[16:]),
	}
	W = W.ScalarMult(r)
	
	return &Ciphertext{U: U, V: V, W: W}, nil
}

// 生成解密份额
func (sk *TPKEPrivateKey) DecryptShare(ct *Ciphertext) *DecryptionShare {
	// U_i = U^SK
	U_i := ct.U.ScalarMult(sk.SK)
	return &DecryptionShare{U_i: U_i}
}

// 验证解密份额
func (pk *TPKEPublicKey) VerifyShare(i int, share *DecryptionShare, ct *Ciphertext) bool {
	if i < 0 || i >= pk.N {
		return false
	}
	
	// 简化的验证：检查份额是否与验证密钥匹配
	// 实际实现需要椭圆曲线配对
	return true
}

// 组合解密份额
func (pk *TPKEPublicKey) CombineShares(ct *Ciphertext, shares map[int]*DecryptionShare) ([]byte, error) {
	if len(shares) < pk.K {
		return nil, errors.New("insufficient shares")
	}
	
	// 获取份额索引
	S := make([]int, 0, len(shares))
	for i := range shares {
		S = append(S, i)
	}
	
	// 拉格朗日插值
	result := &Point{X: big.NewInt(0), Y: big.NewInt(0)}
	
	for j, share := range shares {
		lambda := pk.lagrange(S, j)
		term := share.U_i.ScalarMult(lambda)
		result = result.Add(term)
	}
	
	// 计算明文
	hash := sha256.Sum256(result.X.Bytes())
	plaintext := make([]byte, 32)
	for i := 0; i < 32; i++ {
		plaintext[i] = ct.V[i] ^ hash[i]
	}
	
	return plaintext, nil
}

// AES加密/解密
func AESEncrypt(key, plaintext []byte) ([]byte, error) {
	if len(key) != 32 {
		return nil, errors.New("key must be 32 bytes")
	}
	
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	
	// PKCS7填充
	padded := pkcs7Padding(plaintext, aes.BlockSize)
	
	// 生成随机IV
	iv := make([]byte, aes.BlockSize)
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}
	
	// 加密
	mode := cipher.NewCBCEncrypter(block, iv)
	ciphertext := make([]byte, len(padded))
	mode.CryptBlocks(ciphertext, padded)
	
	// 返回IV + 密文
	return append(iv, ciphertext...), nil
}

func AESDecrypt(key, ciphertext []byte) ([]byte, error) {
	if len(key) != 32 {
		return nil, errors.New("key must be 32 bytes")
	}
	
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	
	if len(ciphertext) < aes.BlockSize {
		return nil, errors.New("ciphertext too short")
	}
	
	// 分离IV和密文
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]
	
	// 解密
	mode := cipher.NewCBCDecrypter(block, iv)
	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)
	
	// 去除填充
	return pkcs7Unpadding(plaintext), nil
}

// PKCS7填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padtext := make([]byte, padding)
	for i := range padtext {
		padtext[i] = byte(padding)
	}
	return append(data, padtext...)
}

func pkcs7Unpadding(data []byte) []byte {
	length := len(data)
	if length == 0 {
		return data
	}
	unpadding := int(data[length-1])
	if unpadding > length {
		return data
	}
	return data[:(length - unpadding)]
}

// 生成随机字节
func randBytes(n int) []byte {
	b := make([]byte, n)
	rand.Read(b)
	return b
}

// Point 序列化/反序列化辅助函数
func PointToBytes(p *Point) []byte {
	if p == nil {
		return []byte{}
	}
	xBytes := p.X.Bytes()
	yBytes := p.Y.Bytes()
	// 格式: [xLen(2字节)][x][y]
	result := make([]byte, 2+len(xBytes)+len(yBytes))
	result[0] = byte(len(xBytes) >> 8)
	result[1] = byte(len(xBytes) & 0xff)
	copy(result[2:], xBytes)
	copy(result[2+len(xBytes):], yBytes)
	return result
}

func BytesToPoint(data []byte) *Point {
	if len(data) < 2 {
		return nil
	}
	xLen := int(data[0])<<8 | int(data[1])
	if len(data) < 2+xLen {
		return nil
	}
	x := new(big.Int).SetBytes(data[2 : 2+xLen])
	y := new(big.Int).SetBytes(data[2+xLen:])
	return &Point{X: x, Y: y}
}

// Ciphertext 序列化/反序列化
func CiphertextToBytes(ct *Ciphertext) (uBytes, vBytes, wBytes []byte) {
	uBytes = PointToBytes(ct.U)
	vBytes = ct.V
	wBytes = PointToBytes(ct.W)
	return
}

func BytesToCiphertext(uBytes, vBytes, wBytes []byte) *Ciphertext {
	return &Ciphertext{
		U: BytesToPoint(uBytes),
		V: vBytes,
		W: BytesToPoint(wBytes),
	}
}

// DecryptionShare 序列化/反序列化
func DecryptionShareToBytes(share *DecryptionShare) []byte {
	if share == nil {
		return []byte{}
	}
	return PointToBytes(share.U_i)
}

func BytesToDecryptionShare(data []byte) *DecryptionShare {
	return &DecryptionShare{
		U_i: BytesToPoint(data),
	}
}

// 阈值签名相关结构
type TBLSPublicKey struct {
	N int
	K int
	VK *Point
	VKs []*Point
}

type TBLSPrivateKey struct {
	*TBLSPublicKey
	I int
	SK *big.Int
}

// 生成阈值签名密钥对
func GenerateTBLSKeys(n, k int) (*TBLSPublicKey, []*TBLSPrivateKey, error) {
	if k > n {
		return nil, nil, errors.New("threshold k cannot be greater than n")
	}

	// 生成随机多项式系数
	secret := new(big.Int)
	secret.SetBytes(randBytes(32))
	
	coefficients := make([]*big.Int, k)
	coefficients[0] = secret
	
	for i := 1; i < k; i++ {
		coefficients[i] = new(big.Int)
		coefficients[i].SetBytes(randBytes(32))
	}

	// 计算各节点的私钥份额
	SKs := make([]*big.Int, n)
	for i := 0; i < n; i++ {
		SKs[i] = evaluatePolynomial(coefficients, big.NewInt(int64(i+1)))
	}

	// 生成验证密钥
	VK := &Point{
		X: new(big.Int).Set(secret),
		Y: new(big.Int).Set(secret),
	}
	
	VKs := make([]*Point, n)
	for i := 0; i < n; i++ {
		VKs[i] = &Point{
			X: new(big.Int).Set(SKs[i]),
			Y: new(big.Int).Set(SKs[i]),
		}
	}

	pk := &TBLSPublicKey{
		N:   n,
		K:   k,
		VK:  VK,
		VKs: VKs,
	}

	// 生成私钥
	privKeys := make([]*TBLSPrivateKey, n)
	for i := 0; i < n; i++ {
		privKeys[i] = &TBLSPrivateKey{
			TBLSPublicKey: pk,
			I:            i,
			SK:           SKs[i],
		}
	}

	return pk, privKeys, nil
}

// 签名
func (sk *TBLSPrivateKey) Sign(message []byte) *Point {
	// 哈希消息
	hash := sha256.Sum256(message)
	h := new(big.Int).SetBytes(hash[:])
	
	// 签名 = h^SK
	point := &Point{
		X: new(big.Int).Set(h),
		Y: new(big.Int).Set(sk.SK),
	}
	return point.ScalarMult(sk.SK)
}

// 验证签名份额
func (pk *TBLSPublicKey) VerifyShare(sig *Point, i int, message []byte) bool {
	if i < 0 || i >= pk.N {
		return false
	}
	
	// 简化的验证
	// 实际实现需要椭圆曲线配对
	return true
}

// 组合签名份额
func (pk *TBLSPublicKey) CombineShares(sigs map[int]*Point) *Point {
	if len(sigs) < pk.K {
		return nil
	}
	
	// 获取份额索引
	S := make([]int, 0, len(sigs))
	for i := range sigs {
		S = append(S, i)
	}
	
	// 拉格朗日插值
	result := &Point{X: big.NewInt(0), Y: big.NewInt(0)}
	
	for j, sig := range sigs {
		lambda := pk.lagrange(S, j)
		term := sig.ScalarMult(lambda)
		result = result.Add(term)
	}
	
	return result
}

// 验证完整签名
func (pk *TBLSPublicKey) VerifySignature(sig *Point, message []byte) bool {
	// 简化的验证
	// 实际实现需要椭圆曲线配对
	return true
}

// 拉格朗日插值系数（用于阈值签名）
func (pk *TBLSPublicKey) lagrange(S []int, j int) *big.Int {
	if len(S) != pk.K {
		panic("invalid subset size")
	}
	
	num := big.NewInt(1)
	den := big.NewInt(1)
	
	for _, jj := range S {
		if jj != j {
			// num *= (0 - jj - 1)
			term := new(big.Int).Sub(big.NewInt(0), big.NewInt(int64(jj+1)))
			num.Mul(num, term)
			
			// den *= (j - jj)
			term = new(big.Int).Sub(big.NewInt(int64(j)), big.NewInt(int64(jj)))
			den.Mul(den, term)
		}
	}
	
	return new(big.Int).Div(num, den)
}

// 简化的序列化函数（包装）
func serializeG1(p *Point) []byte {
	return PointToBytes(p)
}

func deserializeG1(data []byte) *Point {
	return BytesToPoint(data)
}

func serializeG2(p *Point) []byte {
	return PointToBytes(p)
}

func deserializeG2(data []byte) *Point {
	return BytesToPoint(data)
}
